/*
 * Copyright 2022 NVIDIA Corporation.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "VkVideoEncoder/VkEncoderConfig.h"
#include "VkVideoEncoder/VkVideoEncoder.h"
#include "VkCodecUtils/VulkanVideoDisplayQueue.h"
#include "VkCodecUtils/VulkanVideoEncodeDisplayQueue.h"
#include "VkCodecUtils/VulkanEncoderFrameProcessor.h"
#include "VkShell/Shell.h"

#define INPUT_FRAME_BUFFER_SIZE 16

int32_t handle_error(const std::error_code& error)
{
    const auto& errmsg = error.message();
    std::printf("error mapping file: %s, exiting...\n", errmsg.c_str());
    return error.value();
}

int main(int argc, char** argv)
{
    VkSharedBaseObj<EncoderConfig> encoderConfig;
    if (VK_SUCCESS != EncoderConfig::CreateCodecConfig(argc, argv, encoderConfig)) {
        return -1;
    }

    static const char* const requiredInstanceLayerExtensions[] = {
        "VK_LAYER_KHRONOS_validation",
        VK_EXT_DEBUG_REPORT_EXTENSION_NAME,
        nullptr
    };

    static const char* const requiredWsiInstanceExtensions[] = {
        // Required generic WSI extensions
        VK_KHR_SURFACE_EXTENSION_NAME,
        nullptr
    };

    static const char* const requiredDeviceExtension[] = {
#if defined(__linux) || defined(__linux__) || defined(linux)
        VK_KHR_EXTERNAL_MEMORY_FD_EXTENSION_NAME,
        VK_KHR_EXTERNAL_FENCE_FD_EXTENSION_NAME,
#endif
        VK_KHR_SYNCHRONIZATION_2_EXTENSION_NAME,
        VK_KHR_VIDEO_QUEUE_EXTENSION_NAME,
        VK_KHR_VIDEO_ENCODE_QUEUE_EXTENSION_NAME,
        nullptr
    };

    static const char* const requiredWsiDeviceExtension[] = {
        // Add the WSI required device extensions
        VK_KHR_SWAPCHAIN_EXTENSION_NAME,
#if !defined(VK_USE_PLATFORM_WIN32_KHR)
        // VK_EXT_DISPLAY_CONTROL_EXTENSION_NAME,
#endif
        nullptr
    };

    static const char* const optinalDeviceExtension[] = {
        VK_EXT_YCBCR_2PLANE_444_FORMATS_EXTENSION_NAME,
        VK_EXT_DESCRIPTOR_BUFFER_EXTENSION_NAME,
        VK_KHR_BUFFER_DEVICE_ADDRESS_EXTENSION_NAME,
        VK_KHR_PUSH_DESCRIPTOR_EXTENSION_NAME,
        nullptr
    };

    std::vector<const char *> reqInstanceExtensions;
    std::vector<const char *> reqDeviceExtensions;

    if (encoderConfig->validate) {
        reqInstanceExtensions.push_back(VK_EXT_DEBUG_REPORT_EXTENSION_NAME);
    }
    /********** Start WSI instance extensions support *******************************************/
    if (encoderConfig->enableFramePresent) {
        const std::vector<VkExtensionProperties>& wsiRequiredInstanceInstanceExtensions =
                Shell::GetRequiredInstanceExtensions(encoderConfig->enableFrameDirectModePresent);

        for (size_t e = 0; e < wsiRequiredInstanceInstanceExtensions.size(); e++) {
            reqInstanceExtensions.push_back(wsiRequiredInstanceInstanceExtensions[e].extensionName);
        }

        for (uint32_t i = 0; requiredWsiInstanceExtensions[0]; i++) {
            const char* name = requiredWsiInstanceExtensions[i];
            if (name == nullptr) {
                break;
            }
            reqInstanceExtensions.push_back(name);
        }
        // terminate the reqInstanceExtensions list with nullptr
        reqInstanceExtensions.push_back(nullptr);

        // Add the WSI required device extensions
        for (uint32_t i = 0; requiredWsiDeviceExtension[0]; i++) {
            const char* name = requiredWsiDeviceExtension[i];
            if (name == nullptr) {
                break;
            }
            reqDeviceExtensions.push_back(name);
        }
    }
    /********** End WSI instance extensions support *******************************************/

    // Add the Vulkan video required device extensions
    for (uint32_t i = 0; requiredDeviceExtension[0]; i++) {
        const char* name = requiredDeviceExtension[i];
        if (name == nullptr) {
            break;
        }
        reqDeviceExtensions.push_back(name);
    }
    // terminate the reqDeviceExtensions list with nullptr
    reqDeviceExtensions.push_back(nullptr);

    VulkanDeviceContext vkDevCtxt(encoderConfig->deviceId,
            encoderConfig->validate ? requiredInstanceLayerExtensions : nullptr,
            reqInstanceExtensions.data(),
            reqDeviceExtensions.data(),
            optinalDeviceExtension);


    VkResult result = vkDevCtxt.InitVulkanDevice(encoderConfig->appName.c_str(),
                                                 encoderConfig->verbose);
    if (result != VK_SUCCESS) {
        printf("Could not initialize the Vulkan device!\n");
        return -1;
    }

    result = vkDevCtxt.InitDebugReport(encoderConfig->validate,
                                       encoderConfig->validateVerbose);
    if (result != VK_SUCCESS) {
        return -1;
    }

    const bool supportsDisplay = true;
    const int32_t numEncodeQueues = ((encoderConfig->queueId != 0) ||
                                     (encoderConfig->enableHwLoadBalancing != 0)) ?
                                     -1 : // all available HW encoders
                                      1;  // only one HW encoder instance

    VkQueueFlags requestVideoEncodeQueueMask = VK_QUEUE_VIDEO_ENCODE_BIT_KHR;

    VkQueueFlags requestVideoDecodeQueueMask = 0;
    if (encoderConfig->enableVideoDecoder) {
        requestVideoDecodeQueueMask |= VK_QUEUE_VIDEO_DECODE_BIT_KHR |
                                       VK_QUEUE_TRANSFER_BIT;
    }

    if (encoderConfig->selectVideoWithComputeQueue) {
        requestVideoEncodeQueueMask |= VK_QUEUE_COMPUTE_BIT;
        if (encoderConfig->enableVideoDecoder) {
            requestVideoDecodeQueueMask |= VK_QUEUE_COMPUTE_BIT;
        }
    }

    VkSharedBaseObj<VulkanVideoDisplayQueue<VulkanEncoderInputFrame>> videoDispayQueue;
    result = CreateVulkanVideoEncodeDisplayQueue(&vkDevCtxt,
                                                 encoderConfig->input.width,
                                                 encoderConfig->input.height,
                                                 encoderConfig->input.bpp,
                                                 encoderConfig->input.vkFormat,
                                                 videoDispayQueue);
    if (result != VK_SUCCESS) {
        return -1;
    }

    VkSharedBaseObj<VkVideoQueue<VulkanEncoderInputFrame>> videoQueue(videoDispayQueue);
    VkSharedBaseObj<FrameProcessor> frameProcessor;
    result = CreateEncoderFrameProcessor(&vkDevCtxt, videoQueue, frameProcessor);
    if (result != VK_SUCCESS) {
        return -1;
    }

    VkSharedBaseObj<VkVideoEncoder> encoder; // the encoder's instance
    if (supportsDisplay && encoderConfig->enableFramePresent) {

        const Shell::Configuration configuration(encoderConfig->appName.c_str(),
                                                 4, // the display queue size
                                                 encoderConfig->enableFrameDirectModePresent);
        VkSharedBaseObj<Shell> displayShell;
        result = Shell::Create(&vkDevCtxt, configuration, frameProcessor, displayShell);
        if (result != VK_SUCCESS) {
            assert(!"Can't allocate display shell! Out of memory!");
            return -1;
        }

        result = vkDevCtxt.InitPhysicalDevice((VK_QUEUE_GRAPHICS_BIT | requestVideoDecodeQueueMask | requestVideoEncodeQueueMask),
                                               displayShell,
                                               requestVideoDecodeQueueMask,
                                               (VK_VIDEO_CODEC_OPERATION_DECODE_H264_BIT_KHR |
                                                VK_VIDEO_CODEC_OPERATION_DECODE_H265_BIT_KHR),
                                               requestVideoEncodeQueueMask,
                                                VK_VIDEO_CODEC_OPERATION_ENCODE_H264_BIT_KHR |
                                                VK_VIDEO_CODEC_OPERATION_ENCODE_H265_BIT_KHR);
        if (result != VK_SUCCESS) {

            assert(!"Can't initialize the Vulkan physical device!");
            return -1;
        }
        assert(displayShell->PhysDeviceCanPresent(vkDevCtxt.getPhysicalDevice(),
                                                   vkDevCtxt.GetPresentQueueFamilyIdx()));

        result = vkDevCtxt.CreateVulkanDevice(encoderConfig->enableVideoDecoder ? 1 : 0, // num decode queues
                                              numEncodeQueues,   // num encode queues
                                              false,             // createTransferQueue
                                              true,              // createGraphicsQueue
                                              true,              // createDisplayQueue
                                              (encoderConfig->selectVideoWithComputeQueue == 1)  // createComputeQueue
                                              );
        if (result != VK_SUCCESS) {

            assert(!"Failed to create Vulkan device!");
            return -1;
        }

        result = VkVideoEncoder::CreateVideoEncoder(&vkDevCtxt, encoderConfig, encoder);
        if (result != VK_SUCCESS) {
            assert(!"Can't initialize the Vulkan physical device!");
            return -1;
        }

        if (displayShell && videoDispayQueue) {
            result = encoder->AttachDisplayQueue(displayShell, videoDispayQueue);
        }

        // std::this_thread::sleep_for(std::chrono::seconds(5));

    } else {

        // No display presentation and no decoder - just the encoder
        result = vkDevCtxt.InitPhysicalDevice((requestVideoDecodeQueueMask | requestVideoEncodeQueueMask | VK_QUEUE_TRANSFER_BIT),
                                               nullptr,
                                               requestVideoDecodeQueueMask,
                                               (VK_VIDEO_CODEC_OPERATION_DECODE_H264_BIT_KHR |
                                                VK_VIDEO_CODEC_OPERATION_DECODE_H265_BIT_KHR),
                                               requestVideoEncodeQueueMask,
                                                VK_VIDEO_CODEC_OPERATION_ENCODE_H264_BIT_KHR |
                                                VK_VIDEO_CODEC_OPERATION_ENCODE_H265_BIT_KHR);
        if (result != VK_SUCCESS) {

            assert(!"Can't initialize the Vulkan physical device!");
            return -1;
        }

        result = vkDevCtxt.CreateVulkanDevice(encoderConfig->enableVideoDecoder ? 1 : 0, // num decode queues
                                              numEncodeQueues,     // num encode queues
                                              // If no graphics or compute queue is requested, only video queues
                                              // will be created. Not all implementations support transfer on video queues,
                                              // so request a separate transfer queue for such implementations.
                                              ((vkDevCtxt.GetVideoEncodeQueueFlag() & VK_QUEUE_TRANSFER_BIT) == 0), //  createTransferQueue
                                              false, // createGraphicsQueue
                                              false, // createDisplayQueue
                                              (encoderConfig->selectVideoWithComputeQueue == 1)  // createComputeQueue
                                              );
        if (result != VK_SUCCESS) {

            assert(!"Failed to create Vulkan device!");
            return -1;
        }

        result = VkVideoEncoder::CreateVideoEncoder(&vkDevCtxt, encoderConfig, encoder);
        if (result != VK_SUCCESS) {
            assert(!"Can't initialize the Vulkan physical device!");
            return -1;
        }
    }

    // Enter the encoding frame loop
    uint32_t curFrameIndex = 0;
    for(; curFrameIndex < encoderConfig->numFrames; curFrameIndex++) {

        if (encoderConfig->verboseFrameStruct) {
            std::cout << "####################################################################################" << std::endl
                      << "Start processing current input frame index: " << curFrameIndex << std::endl;
        }

        VkSharedBaseObj<VkVideoEncoder::VkVideoEncodeFrameInfo> encodeFrameInfo;
        encoder->GetAvailablePoolNode(encodeFrameInfo);
        assert(encodeFrameInfo);
        // load frame data from the file
        result = encoder->LoadNextFrame(encodeFrameInfo);
        if (result != VK_SUCCESS) {
            std::cout << "ERROR processing input frame index: " << curFrameIndex << std::endl;
            break;
        }

        if (encoderConfig->verboseFrameStruct) {
            std::cout << "End processing current input frame index: " << curFrameIndex << std::endl;
        }
    }

    encoder->WaitForThreadsToComplete();

    std::cout << "Done processing " << curFrameIndex << " input frames!" << std::endl
              << "Encoded file's location is at " << encoderConfig->outputFileHandler.GetFileName()
              << std::endl;
    return 0;
}
